import numpy as np


class QuaternionIntegration:
    def __init__(self) -> None:
        self._q = np.array([1, 0, 0, 0])
        self._ts = 0.01

    def integration(self, omega: np.ndarray):
        eps = 1e-6

        Omega = np.array([
            [0, -omega[0], -omega[1], -omega[2]],
            [omega[0], 0, omega[2], -omega[1]],
            [omega[1], -omega[2], 0, omega[0]],
            [omega[2], omega[1], -omega[0], 0]
        ])

        omega_norm = np.linalg.norm(omega)
        if omega_norm < eps:
            self._q = (np.eye(4) + 0.5 * self._ts * Omega) @ self._q
        else:
            self._q = (np.cos(0.5 * omega_norm * self._ts) * np.eye(4) + 1 / omega_norm * np.sin(
                0.5 * omega_norm * self._ts) * Omega) @ self._q
            # self._q = (np.eye(4) + 0.5 * self._ts * Omega) @ self._q

    @property
    def get_q(self) -> np.ndarray:
        return self._q.copy()
